import numpy
import copy

def get_prev_stim(sample_frame, shade):
    '''
    Return a 2D numpy array of the same dimensions as sample_frame, but
        with all values between 0-255 based on the value of shade between
        -1.0 and 1.0.
    '''
    return_frame = copy.deepcopy(sample_frame)
    frame_val = int(255 * (shade+1.0)/2.0)
    return_frame[:][:] = frame_val
    return return_frame

def stim_chunk(end_time, tau, num_samples, frames, frame_times, shade=1.0):
    """
    This function figures out what was displayed on the screen at a given,
        time and returns a 3 dimensional numpy array whose:
            shape[0] : is of length tau*sample_rate (in approprate units)
            shape[1] : win_size[0]
            shape[2] : win_size[1] where win_size is the size, in pixels
                                   window used to generate the frames.
    Inputs:
        end_time        : The time of the 'event' or spike.
        tau             : The length of time before the 'event' or spike that
                          we're calculating the spike triggered average for.
        num_samples     : len(return_value), number of time samples in chunk.
        frames          : The frames representing the stimulus on the screen.
        frame_times     : the times when the frames were presented.
        shade=1.0       : The window background shade prior to stimulation.
    """
    prev_frame = get_prev_stim(frames[0], shade)
    # ignore the first value of frame_times, since it is always 0.0
    tframe_times = frame_times[1:]
    if end_time < 0.0:
        return None
    if end_time > tframe_times[-1]:
        return None

    begin_time = end_time - tau
    times = numpy.linspace(begin_time, end_time, num_samples)

    for i, ftime in enumerate(tframe_times):
        # find the first time greater than begin_time.
        if ftime > begin_time:
            begin_i = i
            break

    ri = 0
    current_time = times[ri]
    x = frames.shape[1]
    y = frames.shape[2]
    return_value = numpy.empty((num_samples,x,y), dtype=numpy.uint32)
    for i in xrange(begin_i-1,len(tframe_times)-1):
        # if i < 0 add in prev_frame instead of frames[i]
        if i < 0:
            while current_time < 0.0:
                # keep putting in the prev_frame
                return_value[ri] = prev_frame
                ri += 1
                if ri == len(times):
                    return return_value
                current_time = times[ri]
        else:
            while current_time < tframe_times[i+1]:
                # keep putting in the current frame
                return_value[ri] = frames[i]
                ri += 1
                if ri == len(times):
                    return return_value
                current_time = times[ri]
        
def sta(frames, frame_times, spike_times, tau, num_samples, shade=1.0):
    """
    Calculate the spike triggered average.
    Inputs:
        frames      : A 3D numpy array of formated frames. (see format_frames)
        frame_times : A list of times when the frames were drawn.
        spike_times : A list of times when 'events' or spikes occured
        tau         : the window of time BEFORE an 'event' to look at the stimulus.
        num_samples : sets the temporal resolution of the sta.
        shade=1.0       : The window background shade prior to stimulation.
    Returns:
        sta         : A 3D numpy array of formated frames, this time 
                       .shape[0]==num_samples
    """
    if type(frames) == list:
        frames = format_frames(frames)
       
    tau_frames = []
    for st in spike_times:
        tchunk =  stim_chunk(st, tau, num_samples, frames, frame_times) 
        if tchunk is not None:
            tau_frames.append(tchunk)
    
    tau_frames = numpy.array(tau_frames)
    return_value = numpy.average(tau_frames, axis=0)
    return return_value, len(tau_frames)

def view_frame(sta, frame_num):
    import Image
    Image.fromarray(numpy.array(sta,dtype=numpy.uint8)[frame_num].T).show()
    
    

def format_frames(frames):
    """
    This function takes a list of PIL image objects and creates
        a greyscale 3D numpy array (0-255). 
        The PIL image objects have a red, green, and blue, channel
        so the conversion is
        greyscale = (11*R + 16*G + 5*B)/32 as is commonplace.
    """
    import numpy

    frame_1 = numpy.array(frames[0]).T
    x = frame_1.shape[1]
    y = frame_1.shape[2]
    return_value = numpy.empty((len(frames), x, y), dtype=numpy.uint8)

    for i, frame in enumerate(frames):
        nframe = numpy.array(frame, dtype=numpy.uint32).T
        R = nframe[0]
        G = nframe[1]
        B = nframe[2]
        return_value[i] = (R*11 + G*16 + B*5)/32

    return return_value
        
    
